package com.heakey.controller;

import com.heakey.enums.prompt.PromptEnum;
import com.heakey.service.DateTimeTools;
import com.heakey.service.ProductTools;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.core.io.ClassPathResource;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.ollama.api.OllamaOptions;

import java.nio.charset.Charset;
import java.time.LocalDate;

/**
 * @author yuluo
 * @author <a href="mailto:yuluo08290126@gmail.com">yuluo</a>
 */
@CrossOrigin
@RestController
@RequestMapping("/ollama/chat-client")
public class OllamaClientController {

	private static final String DEFAULT_PROMPT = "你好，介绍下你自己！请用中文回答。";

	private final ChatClient ollamaiChatClient;

	public OllamaClientController(ChatModel chatModel, ChatMemory chatMemory) {
		ClassPathResource resource = new ClassPathResource(PromptEnum.CHINESE_MEDICAL.getPath());
		this.ollamaiChatClient = ChatClient.builder(chatModel)
				.defaultSystem(resource, Charset.defaultCharset())
				.defaultAdvisors(
						new PromptChatMemoryAdvisor(chatMemory)
				).defaultTools(new ProductTools())
				.build();
	}

	/**
	 * ChatClient 简单调用
	 */
	@GetMapping("/simple/chat")
	public String simpleChat() {
		return ollamaiChatClient.prompt(DEFAULT_PROMPT).call().content();
	}

	/**
	 * ChatClient 流式调用
	 */
	@GetMapping(value = "/stream/chat",produces = MediaType.TEXT_EVENT_STREAM_VALUE)
	public Flux<String> streamChat(@RequestParam("message")String message) {
		Flux<String> content =  ollamaiChatClient.prompt()
				.user(message)
				.advisors(advisorSpec -> advisorSpec.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY,100))
				.stream()
				.content()
				.map(str -> str.replaceAll("<think>[\\s\\S]*?</think>", ""));
		return  content.concatWith(Flux.just("[complete]"));
	}
}